package com.heeb.djl.boot.controller;

import ai.djl.Application;
import ai.djl.Device;
import ai.djl.ModelException;
import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.output.DetectedObjects;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.training.util.ProgressBar;
import ai.djl.translate.TranslateException;
import ai.djl.translate.Translator;
import com.heeb.djl.boot.util.DetectedObjectUtil;
import com.heeb.djl.boot.util.YoloV5Translator;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RestController;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @Author bjiang
 * @Description //TODO 图片检测
 * @Date 2021/12/31 9:12
 * @Version 1.0
 */
@RestController
@Slf4j
public class ObjectDetectionHbjController {


    @PostMapping("/ObjectDetectionHbj")
    public void ObjectDetection(){
        DetectedObjects detection = null;
        try {
            detection = predict();
        } catch (IOException e) {
            e.printStackTrace();
        } catch (ModelException e) {
            e.printStackTrace();
        } catch (TranslateException e) {
            e.printStackTrace();
        }
        log.info("{}", detection);
    }


    public static void main(String[] args){
        try {
            DetectedObjects detection=predict();
            log.info("{}", detection);
        } catch (IOException e) {
            e.printStackTrace();
        } catch (ModelException e) {
            e.printStackTrace();
        } catch (TranslateException e) {
            e.printStackTrace();
        }
    }
    public static DetectedObjects predict() throws IOException, ModelException, TranslateException {
        //Path imageFile = Paths.get("D:/work/kana1.jpg");
        File file=new File("D:/work/kana1.jpg");
        //Image img = ImageFactory.getInstance().fromFile(imageFile);
        BufferedImage bufImg=ImageIO.read(file);
        Image img = ImageFactory.getInstance().fromImage(bufImg);
        Map<String, Object> arguments = new ConcurrentHashMap<>();
        arguments.put("width", 640);//图片以640宽度进行操作
        arguments.put("height", 640);//图片以640高度进行操作
        arguments.put("resize", true);//调整图片大小
        arguments.put("rescale", true);//图片值编程0-1之间
        //arguments.put("normalize", true);

    /*    arguments.put("toTensor", false);//转换成张量
        arguments.put("range", "0,1");//范围
        arguments.put("normalize", "false");//正态化*/
        //arguments.put("threshold", 0.2);//阈值小于0.2不显示
        //arguments.put("nmsThreshold", 0.5);

        //获取模型分类
        Translator<Image, DetectedObjects> translator = YoloV5Translator.builder(arguments).optSynsetArtifactName("coco.names").build();
        Criteria<Image, DetectedObjects> criteria =
                Criteria.builder()
                        .optApplication(Application.CV.INSTANCE_SEGMENTATION)
                        .setTypes(Image.class, DetectedObjects.class)
                        .optDevice(Device.cpu())
                        .optModelPath(Paths.get("build/input"))
                        .optModelName("yolov5s.torchscript.pt") //获取模型
                        .optTranslator(translator)
                        .optProgress(new ProgressBar())
                        .optEngine("PyTorch")
                        .build();
        try (ZooModel<Image, DetectedObjects> model = criteria.loadModel()) {
            try (Predictor<Image, DetectedObjects> predictor = model.newPredictor()) {
                DetectedObjects detection = predictor.predict(img);
                saveBoundingBoxImage(img, detection);
                return detection;
            }
        }
    }

    /**
     * @Author bjiang
     * @Description //TODO 根据detection绘制图片，输出到 build/output
     * @Date 10:08 2021/12/31
     * @Version 1.0
     * @Param [img, detection]
     * @return void
     */
    private static void saveBoundingBoxImage(Image img, DetectedObjects detection)
            throws IOException {
        Path outputDir = Paths.get("build/output");
        Files.createDirectories(outputDir);
        DetectedObjects detectionNew= DetectedObjectUtil.me().getDetectedObjects(detection);
        img.drawBoundingBoxes(detectionNew);
        Path imagePath = outputDir.resolve("instances2.png");
        img.save(Files.newOutputStream(imagePath), "png");
        System.out.println("Segmentation result image has been saved in"+detectionNew);
    }
}
